###########################################################################
# Created by: CASIA IVA 
# Email: jliu@nlpr.ia.ac.cn
# Copyright (c) 2018
###########################################################################
import torch.nn as nn
from model.DANet.da_att import PAM_Module
from model.DANet.da_att import CAM_Module
from model.resnet import resnet34, resnet50, resnet101, resnet152
import torchsummary


BaseNet_version = "DANet"


class DANet(nn.Module):
    def __init__(self, nclass=1, norm_layer=nn.BatchNorm2d):
        super(DANet, self).__init__()
        self.net = resnet34(pretrained=True)
        self.head = DANetHead(512, nclass, norm_layer)
        self._up_kwargs = {'mode': 'bilinear', 'align_corners': True}

    def forward(self, x):
        imsize = x.size()[2:]
        x = self.net.conv1(x)
        x = self.net.bn1(x)
        x = self.net.relu(x)
        x = self.net.maxpool(x)
        x = self.net.layer1(x)
        x = self.net.layer2(x)
        x = self.net.layer3(x)
        x = self.net.layer4(x)
        x = self.head(x)
        x = nn.functional.interpolate(x, imsize, **self._up_kwargs)
        # print(x.size())
        return x


class DANetHead(nn.Module):
    def __init__(self, in_channels, out_channels, norm_layer):
        super(DANetHead, self).__init__()
        inter_channels = in_channels // 4
        self.conv5a = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())
        
        self.conv5c = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())

        self.sa = PAM_Module(inter_channels)
        self.sc = CAM_Module(inter_channels)
        
        self.conv51 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())
        
        self.conv52 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())

        self.conv8 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1))

    def forward(self, x):
        feat1 = self.conv5a(x)
        sa_feat = self.sa(feat1)
        sa_conv = self.conv51(sa_feat)

        feat2 = self.conv5c(x)
        sc_feat = self.sc(feat2)
        sc_conv = self.conv52(sc_feat)
        
        feat_sum = sa_conv+sc_conv
        sasc_output = self.conv8(feat_sum)
        return sasc_output


if __name__ == '__main__':
    model = DANet()
    model.cuda(0)
    torchsummary.summary(model, (3, 256, 256))

